# Ultralytics YOLO 🚀, GPL-3.0 license
import torch
from ultralytics.yolo.engine.validator import BaseValidator
from ultralytics.yolo.v8.detect import DetectionValidator
from ultralytics.yolo.utils import DEFAULT_CFG, ops
from ultralytics.yolo.utils.metrics import ConfusionMatrix, DetMetrics

from .process import preprocess_batch, postprocess_batch


class GenXValidator(DetectionValidator):
    def __init__(self, dataloader=None, save_dir=None, pbar=None, logger=None, args=None):
        self.ch = 20
        BaseValidator.__init__(self, dataloader, save_dir, pbar, logger, args)
        self.args.task = 'detect'
        self.data_dict = self.args.data
        self.is_coco = False
        self.class_map = None
        self.metrics = DetMetrics(save_dir=self.save_dir)
        self.iouv = torch.linspace(0.5, 0.95, 10)  # iou vector for mAP@0.5:0.95
        self.niou = self.iouv.numel()


    def preprocess(self, batch):
        batch = preprocess_batch(batch, self.device, torch.half if self.args.half else torch.float)

        for k in ["batch_idx", "cls", "bboxes"]:
            batch[k] = batch[k].to(self.device)

        nb = len(batch["img"])
        self.lb = [torch.cat([batch["cls"], batch["bboxes"]], dim=-1)[batch["batch_idx"] == i]
                   for i in range(nb)] if self.args.save_hybrid else []  # for autolabelling

        return batch

    def init_metrics(self, model):
        head = model.model[-1] if self.training else model.model.model[-1]
        val = self.data.get('val', '')  # validation path
        self.class_map = list(range(1000))
        self.args.save_json |= self.is_coco and not self.training  # run on final val if training COCO
        self.nc = head.nc
        self.names = model.names
        self.metrics.names = self.names
        self.metrics.plot = self.args.plots
        self.confusion_matrix = ConfusionMatrix(nc=self.nc)
        self.seen = 0
        self.jdict = []
        self.stats = []

    def postprocess(self, preds, batch = None):
        preds = ops.non_max_suppression(preds,
                                        self.args.conf,
                                        self.args.iou,
                                        labels=self.lb,
                                        multi_label=True,
                                        agnostic=self.args.single_cls,
                                        max_det=self.args.max_det)
        if batch is not None:
            batch = postprocess_batch(batch)
        return preds


def val(cfg=DEFAULT_CFG, use_python=False):
    model = cfg.model or "yolov8n.pt"
    data = cfg.data or "coco128.yaml"

    args = dict(model=model, data=data)
    if use_python:
        from ultralytics import YOLO
        YOLO(model).val(**args)
    else:
        validator = GenXValidator(args=args)
        validator(model=args['model'])


if __name__ == "__main__":
    val()
